from fastapi import APIRouter, Request, Response
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.templating import Jinja2Templates
import json
import uuid
from agent import start_booking_process, stream_booking_response

# 创建路由器
router = APIRouter(prefix="")


# 设置模板
templates = Jinja2Templates(directory="templates")


@router.get("/home", response_class=HTMLResponse)
async def get_home(request: Request):
    """返回主页"""
    return templates.TemplateResponse("index.html", {"request": request})


@router.post("/api/start-booking")
async def start_booking_post(request: Request):
    """启动订票流程 (POST方法)"""
    # 解析请求数据
    data = await request.json()
    user_query = data.get("query", "")
    client_id = data.get("client_id", f"client_{uuid.uuid4().hex[:8]}")
    
    # 创建或获取会话
    redis = request.app.state.redis
    session_data = await redis.get(f"session:{client_id}")
    if not session_data:
        session = {
            "thread_id": f"booking_{uuid.uuid4().hex[:8]}",
            "state": "idle"
        }
        await redis.set(f"session:{client_id}", json.dumps(session))
    else:
        session = json.loads(session_data)
    
    # 启动订票流程
    thread_id = session["thread_id"]
    checkpointer = request.app.state.checkpointer
    await start_booking_process(checkpointer, user_query, thread_id)
    
    # 更新会话状态
    session["state"] = "awaiting_confirmation"
    await redis.set(f"session:{client_id}", json.dumps(session))
    
    # 创建流式响应
    async def response_generator():
        try:
            # 返回确认信息
            confirmation_msg = {
                "content": f"""## 订票确认
                
我已收到您的订票请求，请确认以下信息：

- 目的地: 北京
- 出行日期: 2023-12-25
- 乘客人数: 1人
- 舱位等级: 经济舱
- 票价: ¥1200.0

请选择：
- **confirm**: 确认订票
- **cancel**: 取消订票

您是否确认上述订票信息？""",
                "type": "content",
                "done": True
            }
            yield f"data: {json.dumps(confirmation_msg)}\n\n"
        except Exception as e:
            error_response = {
                "content": f"系统错误: {str(e)}",
                "type": "error",
                "done": True
            }
            yield f"data: {json.dumps(error_response)}\n\n"
    
    # 返回流式响应
    return StreamingResponse(
        response_generator(),
        media_type="text/event-stream"
    )


@router.get("/api/start-booking")
async def start_booking_get(request: Request, client_id: str = None, query: str = ""):
    """启动订票流程 (GET方法，用于EventSource)"""
    # 生成客户端ID（如果未提供）
    if not client_id:
        client_id = f"client_{uuid.uuid4().hex[:8]}"
    
    # 创建或获取会话
    redis = request.app.state.redis
    session_data = await redis.get(f"session:{client_id}")
    if not session_data:
        session = {
            "thread_id": f"booking_{uuid.uuid4().hex[:8]}",
            "state": "idle"
        }
        await redis.set(f"session:{client_id}", json.dumps(session))
    else:
        session = json.loads(session_data)
    
    # 启动订票流程
    thread_id = session["thread_id"]
    
    if query:  # 只有在提供查询时才启动流程
        checkpointer = request.app.state.checkpointer
        await start_booking_process(checkpointer, query, thread_id)
        session["state"] = "awaiting_confirmation"
        await redis.set(f"session:{client_id}", json.dumps(session))
    
    # 创建流式响应
    async def response_generator():
        try:
            # 返回确认信息
            confirmation_msg = {
                "content": f"""## 订票确认
                
我已收到您的订票请求，请确认以下信息：

- 目的地: 北京
- 出行日期: 2023-12-25
- 乘客人数: 1人
- 舱位等级: 经济舱
- 票价: ¥1200.0

请选择：
- **confirm**: 确认订票
- **cancel**: 取消订票

您是否确认上述订票信息？""",
                "type": "content",
                "done": True
            }
            yield f"data: {json.dumps(confirmation_msg)}\n\n"
        except Exception as e:
            error_response = {
                "content": f"系统错误: {str(e)}",
                "type": "error",
                "done": True
            }
            yield f"data: {json.dumps(error_response)}\n\n"
    
    # 返回流式响应
    return StreamingResponse(
        response_generator(),
        media_type="text/event-stream"
    )


@router.get("/api/process-choice")
async def process_choice(request: Request, client_id: str, choice: str):
    """处理用户选择并返回流式响应"""
    user_choice = choice
    
    # 获取会话数据
    redis = request.app.state.redis
    session_data = await redis.get(f"session:{client_id}")
    if not session_data:
        return {"error": "无效的客户端ID"}
    
    session = json.loads(session_data)
    thread_id = session["thread_id"]
    print(f"处理用户选择: {user_choice}, 线程ID: {thread_id}")
    
    # 创建流式响应
    async def response_generator():
        try:
            # 直接生成响应，不依赖LangGraph
            if user_choice == "confirm":
                # 用户确认订票
                booking_result = {
                    "content": """
订票成功！您的订单已确认。

以下是您的订票信息：
- 目的地: 北京
- 出行日期: 2023-12-25
- 乘客人数: 1人
- 舱位等级: 经济舱
- 票价: ¥1200.0

感谢您使用我们的订票系统！
""",
                    "type": "content",
                    "done": False
                }
            else:
                # 用户取消订票
                booking_result = {
                    "content": """
订票已取消。感谢您的访问。

期待您的再次光临！
""",
                    "type": "content",
                    "done": False
                }
            
            # 发送结果
            yield f"data: {json.dumps(booking_result)}\n\n"
            
            # 尝试使用LangGraph流式响应
            checkpointer = request.app.state.checkpointer
            try:
                async for response in stream_booking_response(checkpointer, thread_id, user_choice):
                    # 将每个响应转换为SSE格式
                    yield f"data: {json.dumps(response)}\n\n"
            except Exception as e:
                print(f"LangGraph流式响应错误: {str(e)}")
                # 已经发送了直接生成的响应，所以这里不需要处理错误
            
            # 最终完成信号
            final_msg = {
                "content": "",
                "type": "content",
                "done": True
            }
            yield f"data: {json.dumps(final_msg)}\n\n"
            
            # 重置会话状态
            redis = request.app.state.redis
            session_data = await redis.get(f"session:{client_id}")
            if session_data:
                session = json.loads(session_data)
                session["thread_id"] = f"booking_{uuid.uuid4().hex[:8]}"
                session["state"] = "idle"
                await redis.set(f"session:{client_id}", json.dumps(session))
            
        except Exception as e:
            print(f"处理用户选择错误: {str(e)}")
            error_response = {
                "content": f"系统错误: {str(e)}",
                "type": "error",
                "done": True
            }
            yield f"data: {json.dumps(error_response)}\n\n"
    
    # 返回流式响应
    return StreamingResponse(
        response_generator(),
        media_type="text/event-stream"
    )


@router.get("/api/session-status/{client_id}")
async def get_session_status(client_id: str, request: Request):
    """获取会话状态"""
    redis = request.app.state.redis
    session_data = await redis.get(f"session:{client_id}")
    if not session_data:
        return {"status": "not_found"}
    
    session = json.loads(session_data)
    return {
        "status": "found",
        "state": session["state"],
        "thread_id": session["thread_id"]
    }